#ifndef TINYNDARRAY_INTERNAL_COMMON_H
#define TINYNDARRAY_INTERNAL_COMMON_H
#include <functional>
#include <sstream>
#include <type_traits>
#include <vector>

#include "tinyndarray.h"

namespace tinyndarray {
std::pair<int, int> CvtToSliceIndexItem(std::initializer_list<int> l);

int ResolveAxis(int axis, size_t ndim, const std::string& name);

Axis ResolveAxis(const Axis& axes, size_t ndim, const std::string& name, bool sort = false,
                 bool sort_order_normal = true);
// 定义元素级二元操作函数类型
using BinaryElementOp = std::function<void(void* out, const void* a, const void* b)>;
using UnaryElementOp = std::function<void(void* out, const void* in)>;
using ScalarElementOp = BinaryElementOp;

template <typename F>
inline void RunParallel(int size, F op) {
    // 决定并行参数
    int n_workers = -1, n_batch = -1, batch_size = -1;
    GetParallelParams(size, n_workers, n_batch, batch_size);

    if (n_workers <= 1) {
        // 单线程执行
        for (int i = 0; i < size; i++) {
            op(i);
        }
    } else {
        // 并行执行
        std::atomic<int> next_batch(0);
        std::vector<std::thread> workers;
        workers.reserve(n_workers);

        for (int w = 0; w < n_workers; w++) {
            workers.emplace_back([=, &next_batch, &op]() {
                while (true) {
                    int batch = next_batch++;
                    if (batch >= n_batch) break;

                    int start = batch * batch_size;
                    int end = std::min(start + batch_size, size);

                    for (int i = start; i < end; i++) {
                        op(i);
                    }
                }
            });
        }

        for (auto& worker : workers) {
            worker.join();
        }
    }
}

void ApplyOpSimple(NdArray& ret, UnaryElementOp op_func);

// 一元操作
void ApplyOpSimple(NdArray& ret, const NdArray& src, UnaryElementOp op_func);

// 二元操作
void ApplyOpSimple(NdArray& ret, const NdArray& lhs, const NdArray& rhs, BinaryElementOp op_func,
                   bool check_dtype = true);
// 标量操作（模板类型）
template <typename T>
inline void ApplyOpSimple(NdArray& ret, const NdArray& lhs, T scalar, ScalarElementOp op_func) {
    if (ret.shape() != lhs.shape()) {
        throw std::runtime_error("Shapes must match for scalar operation");
    }

    Dtype dtype = ret.dtype();
    if (dtype != lhs.dtype()) {
        throw std::runtime_error("Data types must match for scalar operation");
    }

    size_t type_size = get_type_size(dtype);
    char* ret_data = static_cast<char*>(ret.data_ptr());
    const char* l_data = static_cast<const char*>(lhs.data_ptr());
    size_t size = ret.size();

    // 将标量转换为void指针
    T scalar_value = scalar;
    const void* scalar_ptr = &scalar_value;

    RunParallel(static_cast<int>(size), [&](int i) {
        op_func(static_cast<void*>(ret_data + i * type_size), static_cast<const void*>(l_data + i * type_size),
                scalar_ptr);
    });
}
}  // namespace tinyndarray
#endif